import random
import re
from typing import List

import pandas as pd
from llama_index.core.schema import TextNode
from tqdm import tqdm

from iflytech_assistant.assistant.dataclasses import LianaiRagData
from iflytech_assistant.es import index

INDEX_NAME = "lianai_examples_20241022"


def load_mapping():
    target_mapping_df = pd.read_excel(
        ".vscode/生成偏好对应关系表.xlsx", sheet_name="细分场景id"
    )
    tags_mapping_df = pd.read_excel(
        ".vscode/生成偏好对应关系表.xlsx", sheet_name="生成偏好对应关系"
    )
    target_mapping = {}
    for i, row in target_mapping_df.iterrows():
        target_mapping[row["id"]] = row["名称"].replace("发", "").strip()

    tags_mapping = {}

    for i, row in tags_mapping_df.iloc[1:].iterrows():
        for j in range(0, len(row), 2):
            if pd.isna(row.iloc[j]):
                continue
            tags_mapping[row.iloc[j + 1]] = {
                "target": row.keys()[j].replace("发", "").strip(),
                "tags": re.sub(r"[^\u4e00-\u9fa5a-zA-Z]", "", row.iloc[j]),
            }

    return target_mapping, tags_mapping


target_mapping, tags_mapping = load_mapping()
lianai_tags = [
    item["tags"] for item in tags_mapping.values() if item["target"] == "恋爱军师"
]


def parse_five_suggestions(content: str) -> List[str]:
    obj = re.match(r"1[.](.*?)2[.](.*?)3[.](.*?)4[.](.*?)5[.](.*?)(A.*|$)", content)
    return [obj.group(i).strip() for i in range(1, 6)]


df = pd.read_excel(".vscode/恋爱助手/viplianai_2024_10_22.xlsx")

# filter out i_clickword is NaN
df = df.dropna(subset=["d_segid"])

result_dict = {
    "userinput": [],
    "stage": [],
    "tag": [],
    "content": [],
    "gender": [],
}
failed = 0
for i, row in tqdm(df.iterrows()):
    userinput = row["userinput"]
    content = row["content"]
    if pd.isna(row["gender"]):
        gender = "default"
    else:
        gender = row["gender"]

    try:
        five_suggestions = parse_five_suggestions(content)
    except:
        failed += 1
        print(f"failed to parse content")
        continue

    try:
        selected = int(row["d_segid"]) - 1
    except:
        print(f"failed to parse row['d_segid']")
        failed += 1
        continue

    content = five_suggestions[selected]
    tag = random.choice(lianai_tags)
    stage = random.choice(["初识期", "热恋期", "表白期", "暧昧期", "稳定期"])

    result_dict["userinput"].append(userinput)
    # 恋爱阶段
    result_dict["stage"].append(stage)
    # 用户偏好
    result_dict["tag"].append(tag)
    # 生成的回复
    result_dict["content"].append(content)
    # 性别
    result_dict["gender"].append(gender)

print(f"parsed {len(df) - failed} rows")
print(f"failed to parse {failed} rows")
df = pd.DataFrame(result_dict)

df = (
    df.groupby(["userinput", "stage", "tag", "gender"])
    .agg(lambda x: "\n".join(x))
    .reset_index()
)
print(f"after groupby, {len(df)} rows")

nodes: List[TextNode] = []

for i, row in df.iterrows():
    data: LianaiRagData = LianaiRagData(
        input=row["userinput"],
        gender=row["gender"],
        stage=row["stage"],
        tag=row["tag"],
        examples=row["content"].split("\n"),
    )
    node: TextNode = data.to_text_node()
    nodes.append(node)

index(nodes, INDEX_NAME)
